{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Invertible Neural Network Tutorial\n",
    "\n",
    "### Model to use: Real-NVP (https://arxiv.org/abs/1605.08803)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environment setup and load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "%matplotlib inline\n",
    "\n",
    "# Do not use GPU\n",
    "os.environ['CUDA_VISIBLE_DEVICES']=''\n",
    "\n",
    "'''\n",
    "\n",
    "(X_train, Y_train), (X_test, Y_test) = tf.keras.datasets.mnist.load_data(path='mnist.npz')\n",
    "\n",
    "print(X_train.shape)\n",
    "print(Y_train.shape)\n",
    "\n",
    "plt.imshow(X_train[0])\n",
    "plt.show()\n",
    "'''\n",
    "\n",
    "dataS = pd.read_csv('Dataset/PCC/MLDataset_quat.csv')\n",
    "dataS = dataS.drop('Unnamed: 0',axis=1)\n",
    "X = dataS.iloc[:,:7].values\n",
    "Y = dataS.iloc[:,7:].values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invertible Neural Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define network for calculating s and t\n",
    "def dense_resnet(inputs, mid_channels, output_channels, num_blocks):\n",
    "    assert len(inputs.shape) == 2\n",
    "    def _blocks(_x, name):\n",
    "        shortcut = _x\n",
    "        _x = tf.layers.dense(_x, mid_channels, activation='relu', name=name+'1')\n",
    "        _x = tf.layers.dense(_x, mid_channels, activation='relu', name=name+'2')\n",
    "        return _x + shortcut\n",
    "    \n",
    "    inputs = tf.layers.dense(inputs, mid_channels, activation='relu', name='initial')\n",
    "    \n",
    "    for i in range(num_blocks):\n",
    "        inputs = _blocks(inputs, '{}'.format(i))\n",
    "    inputs = tf.layers.dense(inputs, mid_channels, activation=None, name='final')\n",
    "    return inputs\n",
    "        \n",
    "#Define mask\n",
    "def get_mask(inputs, reverse_mask, data_format='NHWC', dtype=tf.float32):\n",
    "    shape = inputs.get_shape().as_list()\n",
    "    if len(shape) == 2:\n",
    "        N = shape[-1]\n",
    "        range_n = tf.range(N)\n",
    "        odd_ind = tf.mod(range_n, 2)\n",
    "        \n",
    "        odd_ind = tf.reshape(odd_ind, [-1, N])\n",
    "        checker = odd_ind\n",
    "        \n",
    "    \n",
    "    elif len(shape) == 4:\n",
    "        H = shape[2] if data_format == 'NCHW' else shape[1]\n",
    "        W = shape[3] if data_format == 'NCHW' else shape[2]\n",
    "               \n",
    "        range_h = tf.range(H)\n",
    "        range_w = tf.range(W)\n",
    "        \n",
    "        odd_ind_h = tf.cast(tf.mod(range_h, 2), dtype=tf.bool)\n",
    "        odd_ind_w = tf.cast(tf.mod(range_w, 2), dtype=tf.bool)\n",
    "        \n",
    "        odd_h = tf.tile(tf.expand_dims(odd_ind_h, -1), [1, W])\n",
    "        odd_w = tf.tile(tf.expand_dims(odd_ind_w,  0), [H, 1])\n",
    "                \n",
    "        checker = tf.logical_xor(odd_h, odd_w)\n",
    "        \n",
    "        reshape = [-1, 1, H, W] if data_format == 'NCHW' else [-1, H, W, 1]\n",
    "        checker = tf.reshape(checker, reshape)\n",
    "        \n",
    "    \n",
    "    else:\n",
    "        raise ValueError('Invalid tensor shape. Dimension of the tensor shape must be '\n",
    "                         '2 (NxD) or 4 (NxCxHxW or NxHxWxC), got {}.'.format(inputs.get_shape().as_list()))\n",
    "        \n",
    "        \n",
    "    if checker.dtype != dtype:\n",
    "        checker = tf.cast(checker, dtype)\n",
    "        \n",
    "    if reverse_mask:\n",
    "        checker = 1. - checker\n",
    "        \n",
    "    return checker\n",
    "\n",
    "# Define coupling layer\n",
    "def coupling_layer(inputs, mid_channels, num_blocks, reverse_mask, name='coupling_layer', backward=False, reuse=None):\n",
    "    mask = get_mask(inputs, reverse_mask)\n",
    "    with tf.variable_scope(name) as scope:\n",
    "        if reuse:\n",
    "            scope.reuse_variables()\n",
    "            \n",
    "        if backward:\n",
    "            v1 = inputs * mask\n",
    "            v2 = inputs * (1-mask)\n",
    "            with tf.variable_scope('st1'):\n",
    "                st1 = dense_resnet(inputs=v1, mid_channels=mid_channels, output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3)\n",
    "                s1 = st1[:, 0:tf.shape(inputs)[1]]\n",
    "                rescale1 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.))\n",
    "                s1 = rescale1 * tf.nn.tanh(s1)\n",
    "                t1 = st1[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2]\n",
    "                \n",
    "            u2 = (1-mask)*(v2 - t1)*tf.exp(-s1)\n",
    "        \n",
    "            with tf.variable_scope('st2'):\n",
    "                st2 = dense_resnet(inputs=u2, mid_channels=mid_channels, output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3)\n",
    "                s2 = st2[:, 0:tf.shape(inputs)[1]]\n",
    "                rescale2 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.))\n",
    "                s2 = rescale2 * tf.nn.tanh(s2)\n",
    "                t2 = st2[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2]\n",
    "                \n",
    "            u1 = mask * (v1 - t2)*tf.exp(-s2)\n",
    "            inputs = u1 + u2\n",
    "        \n",
    "        else:\n",
    "            u1 = inputs * mask\n",
    "            u2 = inputs * (1-mask)\n",
    "        \n",
    "            with tf.variable_scope('st2'):\n",
    "                st2 = dense_resnet(inputs=u2, mid_channels=mid_channels, output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3)\n",
    "                s2 = st2[:, 0:tf.shape(inputs)[1]]\n",
    "                rescale2 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.))\n",
    "                s2 = rescale2 * tf.nn.tanh(s2)\n",
    "                t2 = st2[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2]\n",
    "        \n",
    "            v1 = mask * (u1 * tf.exp(s2) + t2)\n",
    "            \n",
    "            with tf.variable_scope('st1'):\n",
    "                st1 = dense_resnet(inputs=v1, mid_channels=mid_channels, output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3)\n",
    "                s1 = st1[:, 0:tf.shape(inputs)[1]]\n",
    "                rescale1 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.))\n",
    "                s1 = rescale1 * tf.nn.tanh(s1)\n",
    "                t1 = st1[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2]\n",
    "        \n",
    "        \n",
    "            v2 = (1-mask) * (u2 * tf.exp(s1) + t1)\n",
    "            inputs = v1 + v2\n",
    "        \n",
    "        return inputs\n",
    "    \n",
    "# Code from https://github.com/chrischute/real-nvp\n",
    "def preprocess(x):\n",
    "    data_constraint = 0.9\n",
    "    y = (x*255. + tf.random.uniform(tf.shape(x), 0, 1))/256.\n",
    "    y = (2 * y - 1) * data_constraint\n",
    "    y = (y + 1) / 2\n",
    "    y = tf.log(y) - tf.log(1-y)\n",
    "    \n",
    "    ldj = tf.nn.softplus(y) + tf.nn.softplus(-y) - tf.nn.softplus(tf.log(1-data_constraint) - tf.log(data_constraint))\n",
    "    sldj = tf.reduce_sum(tf.reshape(ldj, [tf.shape(ldj)[0], -1]), axis=-1)\n",
    "    return y, sldj\n",
    "\n",
    "    \n",
    "def real_nvp(inputs, mid_channels, backward=False, reuse=False):\n",
    "#    \n",
    "    x = inputs\n",
    "    if backward:\n",
    "       \n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=True, name='c4', backward=backward, reuse=reuse)\n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=False, name='c3', backward=backward, reuse=reuse)\n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=True, name='c2', backward=backward, reuse=reuse)\n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=False, name='c1', backward=backward, reuse=reuse)\n",
    "    else:\n",
    "        \n",
    "#        x, sldj = preprocess(inputs)\n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=False, name='c1', backward=backward, reuse=reuse)\n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=True, name='c2', backward=backward, reuse=reuse)\n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=False, name='c3', backward=backward, reuse=reuse)\n",
    "        x = coupling_layer(x, mid_channels, 4, reverse_mask=True, name='c4', backward=backward, reuse=reuse)\n",
    "#    x = tf.nn.sigmoid(x)\n",
    "    return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-2-288fad4ed044>:10: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dense instead.\n",
      "WARNING:tensorflow:From /home/esoc/.local/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n",
      "Original: \n",
      " [[-0.39975804 -0.03595382  0.10423118  2.5258074 ]\n",
      " [-0.64843106 -1.3048775   0.8708211   0.03910542]]\n",
      "Forward: \n",
      " [[  2.719318     0.14667463  -0.80206853   4.1788845 ]\n",
      " [  6.4994426  -10.969748    -3.5650327   -3.517817  ]]\n",
      "Restored: \n",
      " [[-0.39975777 -0.03595387  0.10423115  2.5258079 ]\n",
      " [-0.64843154 -1.3048768   0.8708212   0.03910562]]\n"
     ]
    }
   ],
   "source": [
    "forward_inputs = tf.placeholder(tf.float32, [None, 4])\n",
    "output = real_nvp(forward_inputs, 8, backward=False, reuse=False)\n",
    "\n",
    "backward_inputs = tf.placeholder(tf.float32, [None, 4])\n",
    "#x_restored = real_nvp(backward_inputs, 8, backward=True, reuse=True)\n",
    "x_restored = real_nvp(output, 8, backward=True, reuse=True)\n",
    "\n",
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "X_toy, _ = sess.run(preprocess(X_toy))\n",
    "\n",
    "output_ = sess.run(output, feed_dict={forward_inputs:X_toy})\n",
    "#print(sess.run(preprocess(forward_inputs), feed_dict={forward_inputs:X_toy}))\n",
    "#restored_ = sess.run(x_restored, feed_dict={backward_inputs:output_[0]})\n",
    "restored_ = sess.run(x_restored, feed_dict={forward_inputs:X_toy})\n",
    "print('Original: \\n',X_toy)\n",
    "print('Forward: \\n', output_)\n",
    "print('Restored: \\n', restored_)\n",
    "\n",
    "\n",
    "#restored_ = sess.run(restored)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
